import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

# pip install sentence-transformers
# pip install huggingface_hub[hf_transfer]


# 示例文本数据
documents = [
    "猫喜欢晒太阳",
    "狗是人类最好的朋友",
    "猫和狗是常见的宠物",
    "宠物给人带来快乐"
]
# -------------------------------
# 方法1：使用TF-IDF提取文本特征
# -------------------------------
tfidf_vectorizer = TfidfVectorizer()
tfidf_matrix = tfidf_vectorizer.fit_transform(documents)
# TF-IDF向量化结果
print("TF-IDF向量化结果：")
print(tfidf_matrix.toarray())
# 查看特征词
print("\nTF-IDF特征词：", tfidf_vectorizer.get_feature_names_out())

# -------------------------------
# 方法2：使用Sentence-BERT生成嵌入向量
# -------------------------------
# 加载预训练的Sentence-BERT模型
model = SentenceTransformer('all-MiniLM-L6-v2')
# model = SentenceTransformer('gpt2')
# 生成嵌入向量
sentence_embeddings = model.encode(documents)
# 打印嵌入向量
print("\nSentence-BERT生成的嵌入向量：")
for i, embedding in enumerate(sentence_embeddings):
    print(f"句子{i + 1}的向量：", embedding[:5], "...")  # 打印部分维度，避免输出过长

# -------------------------------
# 比较两种特征提取方法的向量相似性
# -------------------------------
# 计算TF-IDF向量的余弦相似度
tfidf_similarity = cosine_similarity(tfidf_matrix)
print("\nTF-IDF向量之间的余弦相似度：")
print(tfidf_similarity)

# 计算句嵌入向量的余弦相似度
embedding_similarity = cosine_similarity(sentence_embeddings)
print("\n句嵌入向量之间的余弦相似度：")
print(embedding_similarity)


"""
### 代码逐部分解释
这段代码围绕**文本特征提取与相似度计算**展开，通过两种主流方法处理中文文本数据，最终对比两种方法的向量相似性结果，具体逻辑如下：

#### 1. 环境与数据准备
- **依赖导入**：导入`numpy`（数值计算）、`TfidfVectorizer`（TF-IDF特征提取）、`SentenceTransformer`（句嵌入模型）、`cosine_similarity`（余弦相似度计算），并注释了需安装的依赖包（`sentence-transformers`和`huggingface_hub`）。
- **示例数据**：定义`documents`列表，包含4条中文文本（如“猫喜欢晒太阳”“宠物给人带来快乐”），作为后续特征提取的原始数据。


#### 2. 方法1：TF-IDF文本特征提取
- **初始化工具**：创建`TfidfVectorizer`实例，该工具会将文本转换为基于“词频-逆文档频率”（TF-IDF）的数值向量（核心逻辑：高频且在少数文档中出现的词，权重更高，更能代表文本特征）。
- **生成TF-IDF矩阵**：调用`fit_transform(documents)`，对4条文本进行分词、特征词筛选（默认过滤停用词），最终生成`tfidf_matrix`（稀疏矩阵，形状为`[4, 特征词数量]`，每行对应一条文本的TF-IDF向量）。
- **结果输出**：
  - 将稀疏矩阵转为数组打印（直观展示每条文本的TF-IDF向量值）；
  - 调用`get_feature_names_out()`，打印所有被筛选出的特征词（即参与TF-IDF计算的关键词，如“猫”“狗”“宠物”等）。


#### 3. 方法2：Sentence-BERT句嵌入向量生成
- **加载预训练模型**：创建`SentenceTransformer`实例，指定预训练模型`all-MiniLM-L6-v2`（轻量级句嵌入模型，支持多语言，能捕捉文本语义，而非仅依赖单个词的频率）；代码中注释了`gpt2`（因`gpt2`主要用于生成任务，不适合直接生成句嵌入，此处为对比提示，实际未使用）。
- **生成句嵌入向量**：调用`model.encode(documents)`，将每条文本转换为固定维度（`all-MiniLM-L6-v2`默认384维）的稠密向量（核心逻辑：基于Transformer模型的上下文理解，将文本语义压缩为数值向量，相似语义的文本向量更接近）。
- **结果输出**：遍历每条文本的嵌入向量，仅打印前5维（避免384维向量输出过长，同时展示向量格式）。


#### 4. 两种方法的向量相似度对比
- **TF-IDF向量相似度**：调用`cosine_similarity(tfidf_matrix)`，计算4条TF-IDF向量间的两两余弦相似度（取值范围`[-1,1]`，值越接近1表示文本越相似），输出为`4x4`矩阵（矩阵`[i][j]`表示第`i+1`条与第`j+1`条文本的相似度）。
- **句嵌入向量相似度**：调用`cosine_similarity(sentence_embeddings)`，同理计算4条句嵌入向量的两两余弦相似度，输出矩阵格式与TF-IDF结果对应，用于对比两种方法的相似度计算差异（例如：TF-IDF可能因“猫”“狗”等词的重叠度判断相似，而句嵌入能基于“宠物”的语义关联判断相似）。


### 一句话总结功能
该代码通过TF-IDF（词频统计）和Sentence-BERT（语义理解）两种方法，分别将4条中文文本转换为数值向量，并计算两种向量的两两余弦相似度，以此对比两种文本特征提取方法的差异。
"""
